#ifndef AST_PRINTER_H
#define AST_PRINTER_H

#include "ASTNode.h"
#include <iostream>
#include <string>

class ASTPrinter : public ASTVisitor {
private:
    int depth;
    
    void printIndent() {
        for (int i = 0; i < depth; ++i) {
            std::cout << "  ";
        }
    }
    
public:
    ASTPrinter() : depth(0) {}
    
    void print(ASTNode& root) {
        depth = 0;
        root.accept(*this);
    }
    
    void visit(ProgramNode& node) override {
        printIndent();
        std::cout << "Program" << std::endl;
        depth++;
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(CompUnitNode& node) override {
        printIndent();
        std::cout << "CompUnit" << std::endl;
        depth++;
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(BasicTypeNode& node) override {
        printIndent();
        std::cout << "BasicType (" << node.typeName << ")" << std::endl;
    }
    
    void visit(ConstDeclNode& node) override {
        printIndent();
        std::cout << "ConstDecl" << std::endl;
        depth++;
        if (node.bType) node.bType->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(VarDeclNode& node) override {
        printIndent();
        std::cout << "VarDecl" << std::endl;
        depth++;
        if (node.bType) node.bType->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(ConstDefNode& node) override {
        printIndent();
        std::cout << "ConstDef (" << node.identifier << ")" << std::endl;
        depth++;
        if (node.initVal) node.initVal->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(VarDefNode& node) override {
        printIndent();
        std::cout << "VarDef (" << node.identifier << ")" << std::endl;
        depth++;
        if (node.initVal) node.initVal->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(FuncDefNode& node) override {
        printIndent();
        std::cout << "FuncDef (" << node.identifier << ")" << std::endl;
        depth++;
        if (node.returnType) node.returnType->accept(*this);
        if (node.params) node.params->accept(*this);
        if (node.body) node.body->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(FuncFParamNode& node) override {
        printIndent();
        std::cout << "FuncFParam (" << node.identifier << ")" << std::endl;
        depth++;
        if (node.paramType) node.paramType->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(BlockNode& node) override {
        printIndent();
        std::cout << "Block" << std::endl;
        depth++;
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(AssignStmtNode& node) override {
        printIndent();
        std::cout << "AssignStmt" << std::endl;
        depth++;
        if (node.lval) node.lval->accept(*this);
        if (node.expr) node.expr->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(ExprStmtNode& node) override {
        printIndent();
        std::cout << "ExprStmt" << std::endl;
        depth++;
        if (node.expr) node.expr->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(ReturnStmtNode& node) override {
        printIndent();
        std::cout << "ReturnStmt" << std::endl;
        depth++;
        if (node.expr) node.expr->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(IfStmtNode& node) override {
        printIndent();
        std::cout << "IfStmt" << std::endl;
        depth++;
        if (node.condition) node.condition->accept(*this);
        if (node.thenStmt) node.thenStmt->accept(*this);
        if (node.elseStmt) node.elseStmt->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(BinaryExprNode& node) override {
        printIndent();
        std::cout << "BinaryExpr (" << node.op << ")" << std::endl;
        depth++;
        if (node.left) node.left->accept(*this);
        if (node.right) node.right->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(UnaryExprNode& node) override {
        printIndent();
        std::cout << "UnaryExpr (" << node.op << ")" << std::endl;
        depth++;
        if (node.operand) node.operand->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(CallExprNode& node) override {
        printIndent();
        std::cout << "CallExpr (" << node.funcName << ")" << std::endl;
        depth++;
        for (auto& arg : node.args) {
            if (arg) arg->accept(*this);
        }
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(IdentExprNode& node) override {
        printIndent();
        std::cout << "IdentExpr (" << node.identifier << ")" << std::endl;
    }
    
    void visit(LiteralExprNode& node) override {
        printIndent();
        std::cout << "LiteralExpr (" << node.literalType << ": " << node.value << ")" << std::endl;
    }
    
    void visit(ParenExprNode& node) override {
        printIndent();
        std::cout << "ParenExpr" << std::endl;
        depth++;
        if (node.expr) node.expr->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(InitValNode& node) override {
        printIndent();
        std::cout << "InitVal" << std::endl;
        depth++;
        if (node.expr) node.expr->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
    
    void visit(ConstInitValNode& node) override {
        printIndent();
        std::cout << "ConstInitVal" << std::endl;
        depth++;
        if (node.expr) node.expr->accept(*this);
        for (auto& child : node.children) {
            if (child) child->accept(*this);
        }
        depth--;
    }
};

#endif 